AWSTemplateFormatVersion: '2010-09-09'
Description: >
  Step Functions Cost Review Remediator.
  Review-only module that inventories AWS Step Functions state machines,
  workflow type, logging, tracing, encryption, tags, versions, aliases, and
  recent execution counts across enabled Regions. This module intentionally does
  not stop executions, delete workflows, or modify Step Functions resources.

Parameters:
  ProjectName:
    Type: String
    Default: step-functions-cost-review
    Description: Short base name used for created resources.

  IncludeDetails:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
    Description: Include detailed Step Functions inventory in Lambda response.

  IncludeTags:
    Type: String
    Default: 'true'
    AllowedValues: ['true', 'false']
    Description: Include tags where supported.

  MaxStateMachinesPerRegion:
    Type: Number
    Default: 200
    MinValue: 1
    Description: Maximum state machines to inspect per Region.

  MaxExecutionsPerStateMachine:
    Type: Number
    Default: 100
    MinValue: 0
    Description: Maximum recent executions to sample per state machine. Set 0 to skip execution sampling.

  HighExecutionSampleWarningThreshold:
    Type: Number
    Default: 50
    MinValue: 1
    Description: Number of sampled executions that triggers a review flag for execution volume.

  ExcludedRegionsCsv:
    Type: String
    Default: ap-northeast-3
    Description: Comma-separated Regions to skip.

  ScheduleEnabled:
    Type: String
    Default: 'false'
    AllowedValues: ['true', 'false']
    Description: Whether to create an enabled EventBridge schedule.

  ScheduleExpression:
    Type: String
    Default: rate(30 days)
    Description: EventBridge schedule expression used when ScheduleEnabled is true.

Conditions:
  IsScheduleEnabled: !Equals [!Ref ScheduleEnabled, 'true']

Resources:
  ReviewRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: StepFunctionsCostReviewPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Sid: RegionDiscovery
                Effect: Allow
                Action:
                  - ec2:DescribeRegions
                Resource: '*'
              - Sid: StepFunctionsRead
                Effect: Allow
                Action:
                  - states:ListStateMachines
                  - states:DescribeStateMachine
                  - states:ListTagsForResource
                  - states:ListExecutions
                  - states:ListStateMachineVersions
                  - states:ListStateMachineAliases
                  - states:ListMapRuns
                Resource: '*'

  ReviewFunction:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: !Ref ProjectName
      Description: Review-only Lambda for Step Functions cost visibility.
      Runtime: python3.12
      Handler: index.handler
      Timeout: 900
      MemorySize: 512
      Role: !GetAtt ReviewRole.Arn
      Environment:
        Variables:
          INCLUDE_DETAILS: !Ref IncludeDetails
          INCLUDE_TAGS: !Ref IncludeTags
          MAX_STATE_MACHINES_PER_REGION: !Ref MaxStateMachinesPerRegion
          MAX_EXECUTIONS_PER_STATE_MACHINE: !Ref MaxExecutionsPerStateMachine
          HIGH_EXECUTION_SAMPLE_WARNING_THRESHOLD: !Ref HighExecutionSampleWarningThreshold
          EXCLUDED_REGIONS_CSV: !Ref ExcludedRegionsCsv
      Code:
        ZipFile: |
          import json
          import os
          from datetime import datetime, timezone

          import boto3
          from botocore.exceptions import ClientError, EndpointConnectionError

          def utc_now():
              return datetime.now(timezone.utc)

          def env_bool(name, default=False):
              return os.environ.get(name, str(default)).strip().lower() in ("1", "true", "yes", "y")

          def parse_bool(value, default=False):
              if value is None:
                  return default
              return str(value).strip().lower() in ("1", "true", "yes", "y")

          def csv_list(value):
              raw = value or ""
              return [x.strip() for x in raw.split(",") if x.strip()]

          def error_record(service, operation, err):
              if isinstance(err, ClientError):
                  return {
                      "service": service,
                      "operation": operation,
                      "error_code": err.response.get("Error", {}).get("Code", "Unknown"),
                      "error_message": err.response.get("Error", {}).get("Message", str(err))
                  }
              return {
                  "service": service,
                  "operation": operation,
                  "error_code": type(err).__name__,
                  "error_message": str(err)
              }

          def is_access_denied(err):
              text = str(err).lower()
              if isinstance(err, ClientError):
                  code = err.response.get("Error", {}).get("Code", "").lower()
                  msg = err.response.get("Error", {}).get("Message", "").lower()
                  text = f"{code} {msg}"
              return "accessdenied" in text or "not authorized" in text or "explicit deny" in text or "service control policy" in text

          def is_not_supported_or_not_found(err):
              if isinstance(err, ClientError):
                  code = err.response.get("Error", {}).get("Code", "").lower()
                  msg = err.response.get("Error", {}).get("Message", "").lower()
                  text = f"{code} {msg}"
                  return (
                      "statemachinedoesnotexist" in text or
                      "executiondoesnotexist" in text or
                      "validationexception" in text or
                      "unsupportedoperation" in text
                  )
              return False

          def safe_call(region_result, service_name, operation_name, func, default):
              try:
                  return func()
              except EndpointConnectionError as e:
                  region_result["failed"] += 1
                  region_result["failures"].append(error_record(service_name, operation_name, e))
                  return default
              except ClientError as e:
                  if is_not_supported_or_not_found(e):
                      return default
                  if is_access_denied(e):
                      region_result["access_denied"] += 1
                  else:
                      region_result["failed"] += 1
                  region_result["failures"].append(error_record(service_name, operation_name, e))
                  return default
              except Exception as e:
                  region_result["failed"] += 1
                  region_result["failures"].append(error_record(service_name, operation_name, e))
                  return default

          def get_enabled_regions(excluded):
              ec2 = boto3.client("ec2", region_name="us-east-1")
              response = ec2.describe_regions(AllRegions=True)
              regions = []
              for r in response.get("Regions", []):
                  status = r.get("OptInStatus")
                  name = r.get("RegionName")
                  if status in (None, "opt-in-not-required", "opted-in") and name not in excluded:
                      regions.append(name)
              return sorted(regions)

          def paginate(client, operation, result_key, kwargs=None, max_items=None):
              kwargs = dict(kwargs or {})
              items = []
              try:
                  paginator = client.get_paginator(operation)
                  for page in paginator.paginate(**kwargs):
                      items.extend(page.get(result_key, []) or [])
                      if max_items is not None and len(items) >= max_items:
                          return items[:max_items]
                  return items
              except Exception as paginator_error:
                  if "Operation cannot be paginated" not in str(paginator_error):
                      raise

              method = getattr(client, operation)
              token = None
              while True:
                  call_kwargs = dict(kwargs)
                  if token:
                      call_kwargs["nextToken"] = token
                  response = method(**call_kwargs)
                  items.extend(response.get(result_key, []) or [])
                  if max_items is not None and len(items) >= max_items:
                      return items[:max_items]
                  token = response.get("nextToken")
                  if not token:
                      break
              return items

          def tags_to_dict(tags):
              out = {}
              for tag in tags or []:
                  key = tag.get("key") or tag.get("Key")
                  if key:
                      out[key] = tag.get("value") or tag.get("Value")
              return out

          def classify_execution_status(status):
              if status == "RUNNING":
                  return "running"
              if status == "SUCCEEDED":
                  return "succeeded"
              if status == "FAILED":
                  return "failed_exec"
              if status == "TIMED_OUT":
                  return "timed_out"
              if status == "ABORTED":
                  return "aborted"
              return "other"

          def describe_state_machine_item(sm, detail, executions, versions, aliases, tags, include_tags):
              logging = detail.get("loggingConfiguration") or {}
              tracing = detail.get("tracingConfiguration") or {}
              encryption = detail.get("encryptionConfiguration") or {}
              status_counts = {}
              for e in executions:
                  status = e.get("status", "UNKNOWN")
                  status_counts[status] = status_counts.get(status, 0) + 1

              item = {
                  "name": sm.get("name") or detail.get("name"),
                  "state_machine_arn": sm.get("stateMachineArn") or detail.get("stateMachineArn"),
                  "type": sm.get("type") or detail.get("type"),
                  "status": detail.get("status"),
                  "creation_date": sm.get("creationDate") or detail.get("creationDate"),
                  "role_arn_present": bool(detail.get("roleArn")),
                  "definition_present": bool(detail.get("definition")),
                  "logging_level": logging.get("level"),
                  "logging_include_execution_data": logging.get("includeExecutionData"),
                  "logging_destinations_count": len(logging.get("destinations", []) or []),
                  "tracing_enabled": tracing.get("enabled"),
                  "encryption_type": encryption.get("type"),
                  "kms_key_id_present": bool(encryption.get("kmsKeyId")),
                  "versions_count": len(versions),
                  "aliases_count": len(aliases),
                  "sampled_executions": len(executions),
                  "sampled_execution_status_counts": status_counts,
                  "latest_sampled_execution_start_date": max([e.get("startDate") for e in executions if e.get("startDate")] or [None])
              }
              if include_tags:
                  item["tags"] = tags or {}
              return item

          def scan_region(region, settings):
              rr = {
                  "region": region,
                  "state_machines": 0,
                  "standard_state_machines": 0,
                  "express_state_machines": 0,
                  "active_state_machines": 0,
                  "state_machines_with_logging": 0,
                  "state_machines_with_logging_all": 0,
                  "state_machines_with_logging_error": 0,
                  "state_machines_with_logging_fatal": 0,
                  "state_machines_with_logging_off": 0,
                  "state_machines_with_tracing": 0,
                  "state_machines_with_kms": 0,
                  "state_machine_versions": 0,
                  "state_machine_aliases": 0,
                  "sampled_executions": 0,
                  "sampled_running_executions": 0,
                  "sampled_succeeded_executions": 0,
                  "sampled_failed_executions": 0,
                  "sampled_timed_out_executions": 0,
                  "sampled_aborted_executions": 0,
                  "state_machines_with_running_executions": 0,
                  "state_machines_with_failed_or_timed_out_samples": 0,
                  "state_machines_over_execution_sample_threshold": 0,
                  "map_runs": 0,
                  "review_flags": [],
                  "state_machine_inventory": [],
                  "access_denied": 0,
                  "failed": 0,
                  "failures": []
              }

              sfn = boto3.client("stepfunctions", region_name=region)

              state_machines = safe_call(
                  rr,
                  "StepFunctions",
                  "ListStateMachines",
                  lambda: paginate(sfn, "list_state_machines", "stateMachines", max_items=settings["max_state_machines_per_region"]),
                  []
              )

              rr["state_machines"] = len(state_machines)

              for sm in state_machines:
                  arn = sm.get("stateMachineArn")
                  detail = {}
                  if arn:
                      detail = safe_call(
                          rr,
                          "StepFunctions",
                          "DescribeStateMachine",
                          lambda a=arn: sfn.describe_state_machine(stateMachineArn=a),
                          {}
                      )

                  sm_type = sm.get("type") or detail.get("type")
                  if sm_type == "STANDARD":
                      rr["standard_state_machines"] += 1
                  elif sm_type == "EXPRESS":
                      rr["express_state_machines"] += 1

                  if detail.get("status") in (None, "ACTIVE"):
                      rr["active_state_machines"] += 1

                  logging = detail.get("loggingConfiguration") or {}
                  logging_level = logging.get("level")
                  if logging_level and logging_level != "OFF":
                      rr["state_machines_with_logging"] += 1
                  if logging_level == "ALL":
                      rr["state_machines_with_logging_all"] += 1
                  elif logging_level == "ERROR":
                      rr["state_machines_with_logging_error"] += 1
                  elif logging_level == "FATAL":
                      rr["state_machines_with_logging_fatal"] += 1
                  elif logging_level == "OFF":
                      rr["state_machines_with_logging_off"] += 1

                  tracing = detail.get("tracingConfiguration") or {}
                  if tracing.get("enabled"):
                      rr["state_machines_with_tracing"] += 1

                  encryption = detail.get("encryptionConfiguration") or {}
                  if encryption.get("kmsKeyId"):
                      rr["state_machines_with_kms"] += 1

                  tags = {}
                  if settings["include_tags"] and arn:
                      tags = safe_call(
                          rr,
                          "StepFunctions",
                          "ListTagsForResource",
                          lambda a=arn: tags_to_dict(sfn.list_tags_for_resource(resourceArn=a).get("tags", [])),
                          {}
                      )

                  versions = []
                  aliases = []
                  if arn:
                      versions = safe_call(
                          rr,
                          "StepFunctions",
                          "ListStateMachineVersions",
                          lambda a=arn: paginate(sfn, "list_state_machine_versions", "stateMachineVersions", {"stateMachineArn": a}),
                          []
                      )
                      aliases = safe_call(
                          rr,
                          "StepFunctions",
                          "ListStateMachineAliases",
                          lambda a=arn: paginate(sfn, "list_state_machine_aliases", "stateMachineAliases", {"stateMachineArn": a}),
                          []
                      )

                  rr["state_machine_versions"] += len(versions)
                  rr["state_machine_aliases"] += len(aliases)

                  executions = []
                  if arn and settings["max_executions_per_state_machine"] > 0:
                      executions = safe_call(
                          rr,
                          "StepFunctions",
                          "ListExecutions",
                          lambda a=arn: paginate(
                              sfn,
                              "list_executions",
                              "executions",
                              {"stateMachineArn": a},
                              max_items=settings["max_executions_per_state_machine"]
                          ),
                          []
                      )

                  rr["sampled_executions"] += len(executions)
                  local_running = 0
                  local_failed_or_timed = 0

                  for ex in executions:
                      cls = classify_execution_status(ex.get("status"))
                      if cls == "running":
                          rr["sampled_running_executions"] += 1
                          local_running += 1
                      elif cls == "succeeded":
                          rr["sampled_succeeded_executions"] += 1
                      elif cls == "failed_exec":
                          rr["sampled_failed_executions"] += 1
                          local_failed_or_timed += 1
                      elif cls == "timed_out":
                          rr["sampled_timed_out_executions"] += 1
                          local_failed_or_timed += 1
                      elif cls == "aborted":
                          rr["sampled_aborted_executions"] += 1

                  if local_running > 0:
                      rr["state_machines_with_running_executions"] += 1
                  if local_failed_or_timed > 0:
                      rr["state_machines_with_failed_or_timed_out_samples"] += 1
                  if len(executions) >= settings["high_execution_sample_warning_threshold"]:
                      rr["state_machines_over_execution_sample_threshold"] += 1

                  # ListMapRuns requires a mapRunArn or executionArn in newer APIs for some workflows.
                  # This module does not enumerate executions deeply enough to derive all map run ARNs.
                  # Keep map_runs reserved for future enhancement.

                  if settings["include_details"]:
                      rr["state_machine_inventory"].append(
                          describe_state_machine_item(sm, detail, executions, versions, aliases, tags, settings["include_tags"])
                      )

              if rr["state_machines"] > 0:
                  rr["review_flags"].append("step_functions_state_machines_present")
              if rr["express_state_machines"] > 0:
                  rr["review_flags"].append("express_workflows_present")
              if rr["state_machines_with_logging_all"] > 0:
                  rr["review_flags"].append("state_machines_with_all_logging_present")
              if rr["state_machines_with_tracing"] > 0:
                  rr["review_flags"].append("state_machines_with_xray_tracing_present")
              if rr["sampled_running_executions"] > 0:
                  rr["review_flags"].append("running_executions_sampled")
              if rr["sampled_failed_executions"] > 0 or rr["sampled_timed_out_executions"] > 0:
                  rr["review_flags"].append("failed_or_timed_out_executions_sampled")
              if rr["state_machines_over_execution_sample_threshold"] > 0:
                  rr["review_flags"].append("high_sampled_execution_count_present")

              return rr

          def handler(event, context):
              event = event or {}
              excluded = set(csv_list(event.get("excluded_regions_csv") or os.environ.get("EXCLUDED_REGIONS_CSV", "ap-northeast-3")))
              settings = {
                  "include_details": parse_bool(event.get("include_details"), env_bool("INCLUDE_DETAILS", True)),
                  "include_tags": parse_bool(event.get("include_tags"), env_bool("INCLUDE_TAGS", True)),
                  "max_state_machines_per_region": int(event.get("max_state_machines_per_region") or os.environ.get("MAX_STATE_MACHINES_PER_REGION", "200")),
                  "max_executions_per_state_machine": int(event.get("max_executions_per_state_machine") or os.environ.get("MAX_EXECUTIONS_PER_STATE_MACHINE", "100")),
                  "high_execution_sample_warning_threshold": int(event.get("high_execution_sample_warning_threshold") or os.environ.get("HIGH_EXECUTION_SAMPLE_WARNING_THRESHOLD", "50"))
              }

              summary = {
                  "started_at": utc_now().isoformat(),
                  "review_only": True,
                  "include_details": settings["include_details"],
                  "include_tags": settings["include_tags"],
                  "excluded_regions": sorted(list(excluded)),
                  "execution_sampling": {
                      "max_executions_per_state_machine": settings["max_executions_per_state_machine"],
                      "high_execution_sample_warning_threshold": settings["high_execution_sample_warning_threshold"]
                  },
                  "totals": {
                      "regions_scanned": 0,
                      "state_machines": 0,
                      "standard_state_machines": 0,
                      "express_state_machines": 0,
                      "active_state_machines": 0,
                      "state_machines_with_logging": 0,
                      "state_machines_with_logging_all": 0,
                      "state_machines_with_logging_error": 0,
                      "state_machines_with_logging_fatal": 0,
                      "state_machines_with_logging_off": 0,
                      "state_machines_with_tracing": 0,
                      "state_machines_with_kms": 0,
                      "state_machine_versions": 0,
                      "state_machine_aliases": 0,
                      "sampled_executions": 0,
                      "sampled_running_executions": 0,
                      "sampled_succeeded_executions": 0,
                      "sampled_failed_executions": 0,
                      "sampled_timed_out_executions": 0,
                      "sampled_aborted_executions": 0,
                      "state_machines_with_running_executions": 0,
                      "state_machines_with_failed_or_timed_out_samples": 0,
                      "state_machines_over_execution_sample_threshold": 0,
                      "map_runs": 0,
                      "regions_with_state_machines": 0,
                      "regions_with_review_flags": 0,
                      "access_denied": 0,
                      "failed": 0
                  },
                  "region_results": [],
                  "failures": []
              }

              try:
                  regions = get_enabled_regions(excluded)
              except Exception as e:
                  summary["totals"]["failed"] += 1
                  summary["failures"].append(error_record("EC2", "DescribeRegions", e))
                  summary["ended_at"] = utc_now().isoformat()
                  print(json.dumps(summary, indent=2, default=str))
                  return json.loads(json.dumps(summary, default=str))

              total_keys = [
                  k for k in summary["totals"].keys()
                  if k not in ("regions_scanned", "regions_with_state_machines", "regions_with_review_flags")
              ]

              for region in regions:
                  rr = scan_region(region, settings)
                  for key in total_keys:
                      summary["totals"][key] += int(rr.get(key, 0) or 0)

                  if rr.get("state_machines", 0) > 0:
                      summary["totals"]["regions_with_state_machines"] += 1
                  if rr.get("review_flags"):
                      summary["totals"]["regions_with_review_flags"] += 1

                  summary["totals"]["regions_scanned"] += 1
                  summary["failures"].extend([dict(f, region=region) for f in rr.get("failures", [])])
                  summary["region_results"].append(rr)

              summary["ended_at"] = utc_now().isoformat()
              print(json.dumps(summary, indent=2, default=str))
              return json.loads(json.dumps(summary, default=str))

  ReviewScheduleRule:
    Type: AWS::Events::Rule
    Condition: IsScheduleEnabled
    Properties:
      Name: !Sub '${ProjectName}-schedule'
      Description: Scheduled run for Step Functions Cost Review Remediator.
      ScheduleExpression: !Ref ScheduleExpression
      State: ENABLED
      Targets:
        - Arn: !GetAtt ReviewFunction.Arn
          Id: StepFunctionsCostReviewTarget
          Input: '{"source":"eventbridge-schedule"}'

  ReviewSchedulePermission:
    Type: AWS::Lambda::Permission
    Condition: IsScheduleEnabled
    Properties:
      FunctionName: !Ref ReviewFunction
      Action: lambda:InvokeFunction
      Principal: events.amazonaws.com
      SourceArn: !GetAtt ReviewScheduleRule.Arn

Outputs:
  ReviewFunctionName:
    Description: Lambda function name.
    Value: !Ref ReviewFunction

  RemediatorFunctionName:
    Description: Lambda function name alias.
    Value: !Ref ReviewFunction

  ReviewFunctionArn:
    Description: Lambda function ARN.
    Value: !GetAtt ReviewFunction.Arn

  ReviewOnly:
    Description: This module is intentionally review-only.
    Value: 'true'
